# @Time : 2021/8/5 10:06
# @Author : Li Kunlun
# @Description : 多输入通道和多输出通道

import utils as d2l
from mxnet import nd


# 1、多输入通道
def corr2d_multi_in(X, K):
    """ 首先沿着X和K的第0维（通道维）遍历。然后使用*将结果列表变成add_n函数的位置参数（positional argument）来进行相加
    :param
        :param X: 输入数组
        :param K: 卷积核
        :return: 二维互相关计算输出
    """
    return nd.add_n(*[d2l.corr2d(x, k) for x, k in zip(X, K)])


X = nd.array([[[0, 1, 2], [3, 4, 5], [6, 7, 8]],
              [[1, 2, 3], [4, 5, 6], [7, 8, 9]]])
K = nd.array([[[0, 1], [2, 3]], [[1, 2], [3, 4]]])

# [[  56.   72.]
#  [ 104.  120.]]
# <NDArray 2x2 @cpu(0)>
print(corr2d_multi_in(X, K))


# 2、多输出通道
def corr2d_multi_in_out(X, K):
    # 对K的第0维遍历，每次同输入X做互相关计算。所有结果使用stack函数合并在一起
    return nd.stack(*[corr2d_multi_in(X, k) for k in K])


K = nd.stack(K, K + 1, K + 2)
# (3, 2, 2, 2)
print(K.shape)
"""
[
    [
      [[ 0.  1.]
     [ 2.  3.]]
      [[ 1.  2.]
       [ 3.  4.]]
    ]
    [
      [[ 1.  2.]
       [ 3.  4.]]
      [[ 2.  3.]
       [ 4.  5.]]
    ]
    [
      [[ 2.  3.]
       [ 4.  5.]]
      [[ 3.  4.]
       [ 5.  6.]]
    ]
]
<NDArray 3x2x2x2 @cpu(0)>
"""
print(K)

"""
对输入数组X与核数组K做互相关运算。此时的输出含有3个通道。
其中第一个通道的结果与之前输入数组X与多输入通道、单输出通道核的计算结果一致。
    [[[  56.   72.]
      [ 104.  120.]]
    
     [[  76.  100.]
      [ 148.  172.]]
    
     [[  96.  128.]
      [ 192.  224.]]]
    <NDArray 3x2x2 @cpu(0)>
"""
print(corr2d_multi_in_out(X, K))


# 3、1*1卷积层
# 使用全连接层中的矩阵乘法来实现 1×1 卷积
def corr2d_multi_in_out_1x1(X, K):
    c_i, h, w = X.shape
    c_o = K.shape[0]
    X = X.reshape((c_i, h * w))
    K = K.reshape((c_o, c_i))
    Y = nd.dot(K, X)  # 全连接层的矩阵乘法
    return Y.reshape((c_o, h, w))


X = nd.random.uniform(shape=(3, 3, 3))
K = nd.random.uniform(shape=(2, 3, 1, 1))

Y1 = corr2d_multi_in_out_1x1(X, K)
Y2 = corr2d_multi_in_out(X, K)

# True
"""
做 1×1 卷积时，以上函数与之前实现的互相关运算函数corr2d_multi_in_out等价
输出：True
"""
print((Y1 - Y2).norm().asscalar() < 1e-6)
